"""
显示模块
"""
import os
import io
import json
import time
import threading
import pygame
from collections import deque
from typing import Any, Dict, List, Optional, Callable
import zipfile

from .live_renderer import LiveRenderer

class DisplayModule:
    def __init__(self, width: int, height: int, backend_interface=None):
        self.width = width
        self.height = height
        self.backend_interface = backend_interface  # 后端接口引用

        # 渲染器
        self.renderer = LiveRenderer(width, height)

        # 字体初始化
        if not pygame.font.get_init():
            pygame.font.init()
        self.font = pygame.font.SysFont(["SimHei", "Arial"], 16)
        self.small_font = pygame.font.SysFont(["SimHei", "Arial"], 12)

        # 显示控制
        self.hud_on = True
        self.target_fps = 30
        self.actual_fps = 0.0

        # 运行状态（仅UI状态）
        self.live_mode = False
        self.playback_mode = False
        self.is_playing = False
        self.playback_speed = 1.0

        # 显示缓冲（用于UI平滑显示）
        self.max_buffer = 5  # 减少缓冲，因为不需要存储游戏逻辑
        self.display_frames = deque(maxlen=self.max_buffer)
        self.frames_displayed = 0
        self.display_fps = 0.0

        # UI统计数据（从后端获取）
        self.stats_last_ts = time.time()
        self.status_msg = "初始化"
        self.last_error = ""
        self.ui_stats = {
            "round": 0,
            "reward": 0.0,
            "score": 0,
            "health": 100,
            "episode": 0,
            "epsilon": 0.1
        }

        # 回放数据（UI回放功能）
        self.replay_frames = []
        self.replay_index = 0
        self.replay_metadata = {}
        self._last_playback_time = 0

        # 捕获功能（UI功能）
        self.capture_frames = []
        self.is_capturing = False

        # 画布
        self._canvas = pygame.Surface((width, height)).convert()

        # UI监听器
        self._on_stats = None
        self._on_message = None

        # 性能统计（UI性能）
        self._fps_samples = deque(maxlen=30)
        self._last_render_time = time.time()

    def set_backend_interface(self, backend_interface):
        """设置后端接口"""
        self.backend_interface = backend_interface

    def set_listeners(self, on_stats=None, on_message=None):
        """设置外部监听器"""
        self._on_stats = on_stats
        self._on_message = on_message

    def _emit_stats(self, stats: Dict[str, Any]):
        """发送统计数据到监听器"""
        if self._on_stats:
            try:
                self._on_stats(stats)
            except Exception as e:
                print(f"统计监听器错误: {e}")

    def _emit_msg(self, msg: str):
        """发送消息到监听器"""
        if msg and self._on_message:
            try:
                self._on_message(msg)
            except Exception as e:
                print(f"消息监听器错误: {e}")

    def update_from_backend(self, game_state: Dict[str, Any], stats: Dict[str, Any] = None):
        """从后端接收状态更新（替代原来的_on_frame_event）"""
        # 如果在回放模式，忽略后端更新
        if self.playback_mode:
            return

        # 更新UI显示数据
        display_frame = {
            "game_state": game_state,
            "stats": stats or {},
            "timestamp": time.time()
        }

        self.display_frames.append(display_frame)
        self.frames_displayed += 1

        # 更新UI统计
        if isinstance(stats, dict):
            self.ui_stats.update({
                "round": stats.get("round", self.ui_stats["round"]),
                "reward": stats.get("reward", self.ui_stats["reward"]),
                "score": stats.get("score", self.ui_stats["score"]),
                "health": stats.get("health", self.ui_stats["health"]),
                "episode": stats.get("episode", self.ui_stats["episode"]),
                "epsilon": stats.get("epsilon", self.ui_stats["epsilon"])
            })

            # 发送规范化统计数据
            normalized_stats = self._normalize_stats(stats)
            if normalized_stats:
                self._emit_stats(normalized_stats)

        # 计算显示帧率
        now = time.time()
        if now - self.stats_last_ts >= 1.0:
            self.display_fps = self.frames_displayed / (now - self.stats_last_ts)
            self.frames_displayed = 0
            self.stats_last_ts = now

    def _normalize_stats(self, stats: Dict[str, Any]) -> Dict[str, Any]:
        """标准化统计数据格式"""
        if not isinstance(stats, dict):
            return {}

        normalized = {}
        mappings = {
            "episode": ["episode", "episodes", "ep"],
            "epsilon": ["epsilon", "eps", "e"],
            "avg_reward": ["avg_reward", "average_reward", "reward_mean", "mean_reward", "reward"],
            "loss": ["loss", "avg_loss", "loss_mean"]
        }

        for key, candidates in mappings.items():
            for candidate in candidates:
                if candidate in stats:
                    normalized[key] = stats[candidate]
                    break

        return normalized

    # =========================
    # 回放功能（UI功能）
    # =========================

    def openPlayback(self, path: str):
        """打开回放文件"""
        try:
            with open(path, 'r', encoding='utf-8') as f:
                replay_data = json.load(f)

            self.replay_frames = replay_data.get("frames", [])
            self.replay_metadata = replay_data.get("metadata", {})

            if not self.replay_frames:
                self.status_msg = "回放文件为空"
                self._emit_msg(self.status_msg)
                return False

            self.playback_mode = True
            self.is_playing = False
            self.replay_index = 0

            # 清空显示缓冲
            self.display_frames.clear()

            # 显示第一帧
            self._show_replay_frame()

            self.status_msg = f"已加载回放: {len(self.replay_frames)} 帧"
            self._emit_msg(self.status_msg)
            return True

        except Exception as e:
            self.status_msg = f"加载回放失败: {e}"
            self._emit_msg(self.status_msg)
            print(f"加载回放错误: {e}")
            return False

    def _show_replay_frame(self):
        """显示当前回放帧"""
        if not self.playback_mode or not self.replay_frames:
            return

        if 0 <= self.replay_index < len(self.replay_frames):
            frame = self.replay_frames[self.replay_index]
            game_state = frame.get("game_state", {})

            frame_data = {
                "game_state": game_state,
                "stats": {
                    "round": game_state.get("round", 0),
                    "score": game_state.get("score", 0),
                    "health": game_state.get("player", {}).get("hp", 100),
                    "episode": game_state.get("episode", 0)
                },
                "timestamp": time.time()
            }

            self.display_frames.clear()
            self.display_frames.append(frame_data)

            # 更新UI统计
            self.ui_stats.update({
                "round": game_state.get("round", 0),
                "score": game_state.get("score", 0),
                "health": game_state.get("player", {}).get("hp", 100),
                "episode": game_state.get("episode", 0)
            })

    def play(self):
        """播放回放"""
        if self.playback_mode and self.replay_frames:
            self.is_playing = True
            self.status_msg = "回放播放中"
            self._emit_msg(self.status_msg)
            self._last_playback_time = time.time()

    def pause(self):
        """暂停回放"""
        self.is_playing = False
        self.status_msg = "回放已暂停"
        self._emit_msg(self.status_msg)

    def step(self, frames: int = 1):
        """步进回放"""
        if self.playback_mode and self.replay_frames:
            self.replay_index = max(0, min(len(self.replay_frames) - 1,
                                          self.replay_index + frames))
            self._show_replay_frame()
            self.status_msg = f"回放步进到 {self.replay_index + 1}/{len(self.replay_frames)}"
            self._emit_msg(self.status_msg)

    def stopPlayback(self):
        """停止回放模式"""
        self.playback_mode = False
        self.is_playing = False
        self.replay_frames = []
        self.replay_index = 0
        self.replay_metadata = {}
        self.status_msg = "回放模式已退出"
        self._emit_msg(self.status_msg)

    def setSpeed(self, speed: float):
        """设置回放速度"""
        self.playback_speed = max(0.1, min(10.0, float(speed)))
        self._emit_msg(f"回放速度: {self.playback_speed}x")

    # =========================
    # 渲染功能（UI核心）
    # =========================

    def _get_current_display_frame(self) -> Optional[Dict[str, Any]]:
        """获取当前要显示的帧"""
        # 回放模式优先
        if self.playback_mode and self.replay_frames:
            # 自动播放逻辑
            if self.is_playing:
                now = time.time()
                if now - self._last_playback_time > (0.1 / self.playback_speed):
                    if self.replay_index < len(self.replay_frames) - 1:
                        self.replay_index += 1
                        self._show_replay_frame()
                    else:
                        self.is_playing = False
                        self.status_msg = "回放播放完成"
                        self._emit_msg(self.status_msg)
                    self._last_playback_time = now

            # 返回当前回放帧
            if self.display_frames:
                return self.display_frames[0]
            return None

        # 实时模式：从后端获取当前状态
        if self.backend_interface and not self.playback_mode:
            try:
                game_state = self.backend_interface.get_game_state()
                if game_state:
                    return {
                        "game_state": game_state,
                        "stats": self.ui_stats,
                        "timestamp": time.time()
                    }
            except Exception as e:
                print(f"获取后端状态失败: {e}")

        # 从显示缓冲获取
        if self.display_frames:
            return self.display_frames[-1]  # 最新帧

        return None

    def render_tick(self, screen: pygame.Surface):
        """渲染一帧（UI渲染）"""
        start_time = time.time()

        # 获取要渲染的帧
        frame = self._get_current_display_frame()

        if frame is not None:
            game_state = frame.get("game_state", {})
            stats = frame.get("stats", {})

            # 渲染游戏画面
            if game_state:
                self._canvas.fill((0, 0, 0))
                self.renderer.render(self._canvas, game_state)

                # 缩放到目标区域
                scaled = pygame.transform.smoothscale(self._canvas, (self.width, self.height))
                screen.blit(scaled, (0, 0))

                # 捕获帧（UI功能）
                if self.is_capturing:
                    self.capture_frames.append(screen.copy())
            else:
                screen.fill((20, 20, 30))  # 深色背景
        else:
            screen.fill((20, 20, 30))
            # 显示无数据提示
            no_data_text = self.font.render("等待数据...", True, (200, 200, 200))
            text_rect = no_data_text.get_rect(center=(self.width//2, self.height//2))
            screen.blit(no_data_text, text_rect)

        # 渲染HUD
        if self.hud_on:
            self._render_hud(screen)

        # 计算渲染FPS
        render_time = time.time() - start_time
        if render_time > 0:
            self._fps_samples.append(1.0 / render_time)
            self.actual_fps = sum(self._fps_samples) / len(self._fps_samples)

    def _render_hud(self, screen: pygame.Surface):
        """渲染HUD信息"""
        panel_width = 450
        panel_height = 160
        panel = pygame.Rect(8, 8, panel_width, panel_height)

        # 半透明背景
        s = pygame.Surface((panel_width, panel_height))
        s.set_alpha(200)
        s.fill((0, 0, 0))
        screen.blit(s, (8, 8))

        # 边框
        pygame.draw.rect(screen, (100, 100, 100), panel, 2, border_radius=5)

        # 渲染文本
        y = 16
        x = 16

        # 模式信息
        if self.playback_mode:
            mode = "回放"
            state = "播放" if self.is_playing else "暂停"
            color = (0, 255, 0) if self.is_playing else (255, 200, 0)
        elif self.backend_interface:
            mode = "实时"
            state = "已连接" if hasattr(self.backend_interface, 'backend') else "未连接"
            color = (0, 255, 0) if state == "已连接" else (255, 100, 100)
        else:
            mode = "离线"
            state = "无后端"
            color = (100, 100, 100)

        texts = [
            (f"模式: {mode}  状态: {state}", color),
            (f"渲染FPS: {self.actual_fps:.1f}  显示FPS: {self.display_fps:.1f}  缓冲: {len(self.display_frames)}", (230, 230, 230)),
            (f"回合: {self.ui_stats['round']}  得分: {self.ui_stats['score']}  生命: {self.ui_stats['health']}", (230, 230, 230)),
            (f"训练轮次: {self.ui_stats['episode']}  探索率: {self.ui_stats['epsilon']:.3f}", (200, 200, 255)),
            (f"状态: {self.status_msg}", (200, 200, 200))
        ]

        if self.last_error:
            texts.append((f"错误: {self.last_error}", (255, 100, 100)))

        if self.playback_mode and self.replay_frames:
            progress = self.replay_index / max(1, len(self.replay_frames) - 1)
            texts.append((f"进度: {self.replay_index + 1}/{len(self.replay_frames)} ({progress*100:.1f}%)", (100, 200, 255)))

        for text, color in texts:
            txt_surf = self.font.render(text, True, color)
            screen.blit(txt_surf, (x, y))
            y += 20

    # =========================
    # UI控制功能
    # =========================

    def startLive(self, backend_interface=None):
        """启动实时显示模式"""
        if backend_interface:
            self.backend_interface = backend_interface

        self.live_mode = True
        self.playback_mode = False
        self.status_msg = "实时显示已启动"
        self._emit_msg(self.status_msg)

    def stopLive(self):
        """停止实时显示模式"""
        self.live_mode = False
        self.status_msg = "实时显示已关闭"
        self._emit_msg(self.status_msg)

    def toggleHUD(self, on: bool):
        """切换HUD显示"""
        self.hud_on = bool(on)
        self._emit_msg(f"HUD {'开启' if self.hud_on else '关闭'}")

    def setTargetFPS(self, fps: int):
        """设置目标帧率"""
        self.target_fps = max(1, min(120, int(fps)))

    # =========================
    # 捕获和导出功能（UI功能）
    # =========================

    def beginCapture(self):
        """开始捕获帧"""
        self.capture_frames.clear()
        self.is_capturing = True
        self._emit_msg("开始捕获帧")

    def endCapture(self):
        """结束捕获帧"""
        self.is_capturing = False
        self._emit_msg(f"捕获结束，共 {len(self.capture_frames)} 帧")

    def exportVideo(self, dst: str, fps: int = 30) -> bool:
        """导出视频或图片序列"""
        try:
            if dst.endswith('.zip'):
                # 导出为图片序列压缩包
                with zipfile.ZipFile(dst, 'w') as zf:
                    for i, frame in enumerate(self.capture_frames):
                        img_data = pygame.image.tostring(frame, 'RGB')
                        zf.writestr(f'frame_{i:05d}.png', img_data)
                self._emit_msg(f"已导出到 {dst}")
                return True
            else:
                self._emit_msg("不支持的导出格式")
                return False
        except Exception as e:
            self._emit_msg(f"导出失败: {e}")
            return False

    # =========================
    # 清理和状态管理
    # =========================

    def clear_error(self):
        """清除错误信息"""
        self.last_error = ""

    def set_status(self, message: str):
        """设置状态消息"""
        self.status_msg = message
        self._emit_msg(message)

    def get_ui_stats(self) -> Dict[str, Any]:
        """获取UI统计信息"""
        return self.ui_stats.copy()

    def update_size(self, width: int, height: int):
        """更新显示尺寸"""
        self.width = width
        self.height = height
        self._canvas = pygame.Surface((width, height)).convert()
        if self.renderer:
            self.renderer.update_size(width, height)